{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate noisy signals\n",
    "\n",
    "The default Langmuir samples are not noisy and, thus, do not properly reflect real world data.  This notebook will add noise to the signals to help simulate real world data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import astropy.units as u\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from plasmapy.utils.decorators import validate_quantities\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = [10.5, 0.56 * 10.5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# noise generator\n",
    "\n",
    "\n",
    "@validate_quantities\n",
    "def add_noise(signal, snr_db: u.dB = 10 * u.dB):\n",
    "    r\"\"\"\n",
    "    Add noise to a signal based on a signal-to-noise ratio (SNR).\n",
    "\n",
    "    .. math::\n",
    "\n",
    "        SNR &= \\\\frac{P_{signal}}{P_{noise}}\n",
    "\n",
    "        SNR_{dB} &= 10 \\\\log_{10}(SNR)\n",
    "\n",
    "    where :math:`P_{signal}` and :math:`P_{noise}` is the average power of\n",
    "    signal and noise components, respectively.\n",
    "\n",
    "    The noise component is generated using `numpy.random.normal`.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    signal: array_like\n",
    "        1D array to add noise to\n",
    "\n",
    "    snr_db: `~astropy.units.Quantity`\n",
    "        desired signal-to-noise ratio in decibels\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    array_like\n",
    "        the original signal with noise added to it\n",
    "\n",
    "    References\n",
    "    ----------\n",
    "    .. [1] https://en.wikipedia.org/wiki/Signal-to-noise_ratio\n",
    "    .. [2] https://stackoverflow.com/a/53688043\n",
    "\n",
    "    \"\"\"\n",
    "    sig_pwr = signal**2\n",
    "    sig_pwr_ave = np.mean(sig_pwr)\n",
    "    sig_pwr_ave_db = 10.0 * np.log10(sig_pwr_ave)\n",
    "\n",
    "    noise_pwr_ave_db = sig_pwr_ave_db - snr_db.value\n",
    "    noise_pwr_ave = 10.0 ** (noise_pwr_ave_db / 10.0)\n",
    "\n",
    "    noise = np.random.normal(0.0, np.sqrt(noise_pwr_ave), signal.size)\n",
    "\n",
    "    return signal + noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in example data\n",
    "# this can be found at\n",
    "#   https://github.com/PlasmaPy/PlasmaPy/blob/main/docs/notebooks/langmuir_samples/Beckers2017.npy\n",
    "#\n",
    "filename = \"Beckers2017\"\n",
    "voltage, current = np.load(filename + \".npy\")\n",
    "\n",
    "\n",
    "# add some artificial noise to simulate reliztic digitized signals\n",
    "v_noisy = add_noise(voltage, 38 * u.dB)\n",
    "i_noisy = add_noise(current, 28 * u.dB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figheight = 2.0 * figheight\n",
    "fig, axs = plt.subplots(3, 1, figsize=[figwidth, figheight])\n",
    "\n",
    "# Current plot\n",
    "axs[0].set_ylabel(\"Current\", fontsize=12)\n",
    "axs[0].set_xlabel(\"Index\", fontsize=12)\n",
    "axs[0].plot(i_noisy, label=\"noisy\")\n",
    "axs[0].plot(current, color=\"r\", label=\"original\")\n",
    "axs[0].legend(fontsize=12)\n",
    "\n",
    "# Voltage plot\n",
    "axs[1].set_ylabel(\"Voltage\", fontsize=12)\n",
    "axs[1].set_xlabel(\"Index\", fontsize=12)\n",
    "axs[1].plot(v_noisy, label=\"noisy\")\n",
    "axs[1].plot(voltage, color=\"r\", label=\"original\")\n",
    "axs[1].legend(fontsize=12)\n",
    "\n",
    "# IV plot\n",
    "axs[2].set_ylabel(\"Current\", fontsize=12)\n",
    "axs[2].set_xlabel(\"Voltage\", fontsize=12)\n",
    "axs[2].plot(v_noisy, i_noisy, label=\"noisy\")\n",
    "axs[2].plot(voltage, current, color=\"r\", label=\"original\")\n",
    "axs[2].legend(fontsize=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save = False\n",
    "if save:\n",
    "    savefile = filename + \"_noisy.npy\"\n",
    "    with open(savefile, \"wb\") as fp:\n",
    "        np.save(fp, [v_noisy, i_noisy])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "plasmapy37-nb",
   "language": "python",
   "name": "plasmapy37-nb"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  },
  "nbsphinx": {
   "orphan": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
